# Name: Aovs
# Author: Wang Chenxu
# Date created: 2023-8-20 14:20:19
# last edited: 2023-8-20 14:20:19
# Description: This is a tool which allows you to create aov beauty restore based compositing with multi pass exr file
# coding: utf-8
import nuke
def aovs():
    [gapX, gapY] = [200, 150]

    tn = nuke.selectedNode()
    tn.setSelected(False)
    channels = tn.channels()
    [oriX, oriY] = [tn.xpos(), tn.ypos()]

    # Get screen size of dot and shuffle node
    tmpShuffle = nuke.createNode('Shuffle', inpanel=False)
    [sWidth, sHeight] = [tmpShuffle.screenWidth(), tmpShuffle.screenHeight()]
    nuke.delete(tmpShuffle)
    tmpDot = nuke.createNode('Dot', inpanel=False)
    [dWidth, dHeight] = [tmpDot.screenWidth(), tmpDot.screenHeight()]
    nuke.delete(tmpDot)

    lightList = []
    layerList = []
    # Get light list
    for channel in channels:
        if 'RGBA_' in channel:
            lightList.append(channel.split('.')[0].replace('RGBA_', ''))
            lightList = list(set(lightList))

    for light in lightList:
        # Get layers
        for channel in channels:
            if not 'RGBA_' in channel:
                if light in channel:
                    layerList.append(channel.split('.')[0])

    layerList = list(set(layerList))
    layerList.sort()

    # layerPerLightAcount is counting how many layers per light have
    # used in create sub merge nodes for each direct and indirect merge node
    layerPerLightAcount = int(len(layerList)/len(lightList))

    # Create nodes`````````````````````````````````````````````````````````````````````````````````````````````
    # Create original nodes
    dotYpos = oriY + gapY * 5                                          # dotYpos is every dot up node's y position
    firstDot = nuke.createNode('Dot', inpanel=False)
    firstDot.setSelected(False)
    firstDot.setInput(0, tn)
    firstDot.setXYpos(int(oriX + 0.5 * (sWidth - dWidth)), dotYpos - gapY * 3)

    unpremult = nuke.createNode('Unpremult', inpanel=False)
    unpremult['channels'].setValue('all')
    unpremult.setSelected(False)
    unpremult.setInput(0, firstDot)
    unpremult.setXYpos(oriX, int(dotYpos - gapY * 2.5))

    dot = nuke.createNode('Dot', inpanel=False)
    dot.setSelected(False)
    dot.setInput(0, unpremult)
    dot.setXYpos(int(oriX + 0.5 * (sWidth - dWidth)), dotYpos)

    # Create light shuffle nodes
    index = 0
    lightCount = 0

    # `````````all nodes below are temp nodes```````````````````````````
    lastDot = ''    # Define the last dot first
    lastMerge = ''   # Define the lastMerge first
    lastSubMerge = ''  # Define the last subMerge first
    lastCloseMerge = ''  # Define the lastCloseMerge first
    # ``````````````````````````````````````````````````````````````````
    for light in lightList:
        for layer in layerList:
            if light in layer:
                if not index == 0:

                    # Dot up
                    d = nuke.createNode('Dot', inpanel=False)
                    if not index == 1:
                        d.setInput(0, lastDot)
                        lastDot = d
                    else:
                        d.setInput(0, dot)
                        lastDot = d
                    d.setSelected(False)
                    d.setXYpos(int(oriX + (index + lightCount) * gapX + 0.5 * (sWidth - dWidth)), dotYpos)

                    # the reason for adding the first shuffle and the rest seperately is they have different inputs
                    # mid layer dot
                    mld = ''    # Define the mld first
                    if index % 2 == 0:
                        mld = nuke.createNode('Dot', inpanel=False)
                        mld['label'].setValue(layer.split('_')[0])
                        mld['note_font_size'].setValue(30)
                        mld.setInput(0, d)
                        mld.setSelected(False)
                        mld.setXYpos(int(oriX + (index + lightCount) * gapX + 0.5 * (sWidth - dWidth)), int(dotYpos + 0.5 * gapY))

                    # Shuffle next
                    s = nuke.createNode('Shuffle', inpanel=False)
                    if index % 2 == 0:
                        s.setInput(0, mld)
                    else:
                        s.setInput(0, d)
                    s['in'].setValue(layer)
                    s['label'].setValue('[value in]')
                    s.setSelected(False)
                    s.setXYpos(oriX + (index + lightCount) * gapX, dotYpos + gapY)

                else:
                    # mid layer dot
                    mld = nuke.createNode('Dot', inpanel=False)
                    mld['label'].setValue(layer.split('_')[0])
                    mld['note_font_size'].setValue(30)
                    mld.setInput(0, dot)
                    mld.setSelected(False)
                    mld.setXYpos(int(oriX + (index + lightCount) * gapX + 0.5 * (sWidth - dWidth)), int(dotYpos + 0.5 * gapY))

                    # Shuffle first
                    s = nuke.createNode('Shuffle', inpanel=False)
                    s.setInput(0, mld)
                    s['in'].setValue(layer)
                    s['label'].setValue('[value in]')
                    s.setSelected(False)
                    s.setXYpos(oriX + (index + lightCount) * gapX, dotYpos + gapY)

                # Merge
                if index % 2 == 0:
                    m = nuke.createNode('Merge2', inpanel=False)
                    m['operation'].setValue('plus')
                    m.setInput(0, s)
                    m.setSelected(False)
                    m.setXYpos(oriX + (index + lightCount) * gapX, dotYpos + gapY * 2)
                    lastMerge = m

                if not index % 2 == 0:
                    # MergeDot
                    md = nuke.createNode('Dot', inpanel=False)
                    md.setSelected(False)
                    md.setInput(0, s)
                    md.setXYpos(int(oriX + (index + lightCount) * gapX + 0.5 * (sWidth - dWidth)), int(dotYpos + gapY * 2 + 0.5 * (sHeight - dWidth)))
                    lastMerge.setInput(1, md)

                if index % 2 == 0:
                    if index % layerPerLightAcount == 0:
                        lastSubMerge = m
                    else:
                        # subMergeDot
                        smd = nuke.createNode('Dot', inpanel=False)                                                                                               #↓ This 2 is for 2 units after merge and shuffle
                        smd.setXYpos(int(oriX + (index + lightCount) * gapX + 0.5 * (sWidth - dWidth)), int(dotYpos + gapY * (int(index % layerPerLightAcount/2) + 2) + 0.5 * (sHeight - dWidth)))
                        smd.setInput(0, m)
                        # subMerge
                        sm = nuke.createNode('Merge2', inpanel=False)
                        sm['operation'].setValue('plus')
                        sm.setInput(0, lastSubMerge)
                        sm.setInput(1, smd)
                        sm.setXYpos(oriX + (int(index / layerPerLightAcount) * layerPerLightAcount + lightCount) * gapX, dotYpos + gapY * (int(index % layerPerLightAcount/2) + 2))
                        lastSubMerge = sm
                        if index % layerPerLightAcount == layerPerLightAcount - 2:
                            if lightCount == 0:
                                lastCloseMerge = sm

                            if lightCount > 0:
                                # closeMergeDot
                                cmd = nuke.createNode('Dot', inpanel=False)
                                cmd.setXYpos(sm.ypos(), lastCloseMerge.ypos() + int(gapY*1.5))
                                cmd.setSelected(False)
                                cmd.setInput(0, sm)
                                cmd.setXYpos(int(oriX + (int(index / layerPerLightAcount) * layerPerLightAcount + lightCount) * gapX + 0.5 * (sWidth-dWidth)), int(lastCloseMerge.ypos() + (gapY * 1.5) + 0.5 * (sHeight-dHeight)))

                                # closeMerge
                                cm = nuke.createNode('Merge2', inpanel=False)
                                cm['operation'].setValue('plus')
                                cm.setSelected(False)
                                cm.setInput(0, lastCloseMerge)
                                cm.setInput(1, cmd)
                                cm.setXYpos(oriX, lastCloseMerge.ypos() + int(gapY*1.5))
                                lastCloseMerge = cm
                # BackDrop
                if index % layerPerLightAcount == 0:
                    b = nuke.createNode('BackdropNode', inpanel=False)
                    b.setXYpos(int(oriX + (index + lightCount) * gapX - gapX * 0.5), dotYpos - 2 * gapY)
                    b['bdwidth'].setValue(gapX * (layerPerLightAcount + 0.5))
                    b['bdheight'].setValue(gapY * (layerPerLightAcount/2 + 3.5))
                    b['label'].setValue(light)
                    b['note_font_size'].setValue(130)


                index += 1
        lightCount += 1

    # Create emission
    for channel in channels:
        if 'emission' in channel:
            # Dot up
            d = nuke.createNode('Dot', inpanel=False)
            d.setInput(0, lastDot)
            d.setSelected(False)
            d.setXYpos(int(oriX + (index + lightCount) * gapX + + 0.5 * (sWidth - dWidth)), dotYpos)

            # Shuffle
            emi = nuke.createNode('Shuffle', inpanel=False)
            emi.setInput(0, d)
            emi['in'].setValue('emission')
            emi['label'].setValue('[value in]')
            emi.setSelected(False)
            emi.setXYpos(oriX + (index + lightCount) * gapX, dotYpos + gapY)

            # closeMergeDot
            cmd = nuke.createNode('Dot', inpanel=False)
            cmd.setXYpos(emi.ypos(), lastCloseMerge.ypos() + int(gapY * 1.5))
            cmd.setSelected(False)
            cmd.setInput(0, emi)
            cmd.setXYpos(int(oriX + (int(index / layerPerLightAcount) * layerPerLightAcount + lightCount) * gapX + 0.5 * (sWidth - dWidth)), int(lastCloseMerge.ypos() + (gapY * 1.5) + 0.5 * (sHeight - dHeight)))


            # closeMerge
            cm = nuke.createNode('Merge2', inpanel=False)
            cm['operation'].setValue('plus')
            cm.setSelected(False)
            cm.setInput(0, lastCloseMerge)
            cm.setInput(1, cmd)
            cm.setXYpos(oriX, lastCloseMerge.ypos() + int(gapY * 1.5))
            lastCloseMerge = cm

            # BackDrop
            if index % layerPerLightAcount == 0:
                b = nuke.createNode('BackdropNode', inpanel=False)
                b.setXYpos(int(oriX + (index + lightCount) * gapX - gapX * 0.5), dotYpos - 2 * gapY)
                b['bdwidth'].setValue(gapX * (layerPerLightAcount + 0.5))
                b['bdheight'].setValue(gapY * (layerPerLightAcount / 2 + 3.5))
                b['label'].setValue('emission')
                b['note_font_size'].setValue(130)
            break

    # Create Copy alpha
    c = nuke.createNode('Copy', inpanel=False)
    c.setSelected(False)
    c.setXYpos(oriX, lastCloseMerge.ypos() + gapY * 2)
    c.setInput(0, lastCloseMerge)

    # Side Up Dot
    sud = nuke.createNode('Dot', inpanel=False)
    sud.setSelected(False)
    sud.setXYpos(firstDot.xpos() - gapX, firstDot.ypos())
    sud.setInput(0, firstDot)

    # Side Down Dot
    sdd = nuke.createNode('Dot', inpanel=False)
    sdd.setSelected(False)
    sdd.setXYpos(int(c.xpos() - gapX + 0.5 * (sWidth - dWidth)), int(c.ypos() + 0.5 * (36 - dHeight))) # 36 is for the hight of Copy node with alpha channel copied
    sdd.setInput(0, sud)

    c.setInput(1, sdd)

    # Create premult for final comp
    p = nuke.createNode('Premult', inpanel=False)
    p.setSelected(False)
    p.setInput(0, c)
    p.setXYpos(oriX, c.ypos() + gapY)